from crewai import Agent
from conn import llms,mysql_conn as mc
from pydantic import BaseModel, Field
import json
from crewai.tools import tool

MC = mc.MysqlConn()

class Entity(BaseModel):
    name: str = Field(..., description="实体名")
    type: str = Field(..., description="实体类型,枚举[A,B,C,D]")
    description: str = Field(..., description="实体描述")

class ResponseModel(BaseModel):
    entitys :list[Entity] = Field(..., description="实体列表")

@tool
def check_entitys_exist(entity_names: list[str]=Field(..., description="实体名列表")):
    """筛选数据库中不存在的实体,返回出来的数据中不存在的实体"""
    entitys = MC.check_entitys_exist(entity_names)
    return [set(entity_names)-set(entitys)]



# 获取页面内容和翻译
def entity_recognition(url: str):
    agent = Agent(
        role="实体识别助手",
        goal='''
        按步骤做事：
        1. 从传入的内容中识别其中的命名实体,命名实体类型有[知名人物名,公司名,产品名,技术名],
        2. 然后从查询一下现有数据库中是否已经存在实体。
        3. 给数据中不存在的实体生成描述并且返回
        返回数据中，类型用枚举[A,B,C,D],分别对应[A:知名人物名,B:公司名,C:产品名,D:技术名]
        ''',
        llm=llms.get_llm(),
        backstory="非常专业的实体识别助手",
        tools=[check_entitys_exist],
        verbose=True,
    )

    r = agent.kickoff(url, response_format=ResponseModel)
    return json.loads(r.raw)['entitys']



if __name__ == "__main__":
    a = '本文详细介绍了跨站请求伪造（CSRF）这一安全漏洞，即攻击者诱骗用户的浏览器向目标网站发送未经授权的请求，利用用户的身份验证Cookie。文章讨论了多种防御措施，如CSRF令牌、SameSite Cookie、Origin头检查和Fetch元数据（Sec-Fetch-Site）。推荐现代应用以Fetch元数据为主要防御手段，并辅以Origin头检查和严格的源验证。文章还强调了保护HTTP→HTTPS CSRF攻击的重要性，并提供了2025年实施CSRF防护的指导，包括Go语言中间件实现。'
    print(entity_recognition(a))
    #a = {'entitys': [{'name': 'CSRF', 'type': '技术名', 'description': '跨站请求伪造（Cross-Site Request Forgery）是一种安全漏洞，攻击者诱骗用户浏览器向目标网站发送未经授权的请求，利用用户的身份验证Cookie。'}, {'name': 'SameSite Cookie', 'type': '技术名', 'description': 'SameSite Cookie 是一种用于防止跨站请求伪造（CSRF）的安全机制，通过限制Cookie在跨站请求中的发送，增强安全性。'}, {'name': 'Origin头检查', 'type': '技术名', 'description': 'Origin头检查是一种用于验证HTTP请求来源的技术，确保请求来自预期的源，从而防范CSRF等攻击。'}, {'name': 'Fetch元数据', 'type': '技术名', 'description': 'Fetch元数据（如Sec-Fetch-Site）是现代浏览器提供的附加信息，帮助服务器判断请求是否来自可信上下文，常用于CSRF防护。'}, {'name': 'Sec-Fetch-Site', 'type': '技术名', 'description': 'Sec-Fetch-Site 是一个HTTP头部字段，用于指示请求的来源上下文，有助于检测和防御跨站请求伪造（CSRF）攻击。'}, {'name': 'Go语言中间件', 'type': '技术名', 'description': 'Go语言中间件是一种在Go语言Web应用中实现的功能模块，可用于添加CSRF防护逻辑，例如生成和验证CSRF令牌。'}]}
    #MC.insert('entitys',a['entitys'])